#ifndef CKDTREE_H
#define CKDTREE_H

#include <vector>
#include <cmath>
#include "CRRTNode.h"
#include "CPoint.h"

class CKDTree
{
public:
    CKDTree(const CKDTree&) = delete;
    CKDTree& operator = (const CKDTree&) = delete;

    template<typename iterator>
    CKDTree(iterator begin, iterator end)   :   m_nodes(begin, end) 
    {
        m_root = makeTree(0, m_nodes.size(), 0);
    }

    template<typename func>
    CKDTree(func& f, size_t n)
    {
        m_nodes.reserve(n);
        for (size_t i = 0; i < n; i++)
        {
            m_nodes.push_back(f(i));
        }

        m_root = makeTree(0, m_nodes.size(), 0);
    }

    bool empty()        const { return m_nodes.empty(); }

    size_t visited()    const { return m_visited;  }

    double distance()   const { return std::sqrt(m_best_distance);  }

    const CPoint2D& nearest(const CPoint2D& pt)
    {
        if (m_root == nullptr)
        {
            throw std::logic_error("kdtree is empty\n");
        }

        m_best = nullptr;

        m_visited = 0;

        m_best_distance = 0;

        nearest(m_root, pt, 0);

        return m_best->_point;
    }

private:

    struct TKDNode
    {
        TKDNode(const CPoint2D& pt)    :   _point(pt), _left(nullptr), _right(nullptr) {}
        int getX()   const { return _point.getX(); }
        int getY()   const { return _point.getY(); }
        int get(size_t index)
        {
            return (index == 0 ? _point.getX() : _point.getY());
        }
        double distance(const CPoint2D& point) { return _point.getSquareDistance(point); }

        CPoint2D _point;
        TKDNode* _left;
        TKDNode* _right;
    };

private:

    struct node_cmp
    {
        node_cmp(size_t index) : _index(index)  {}
        bool operator()(const TKDNode& n1, const TKDNode& n2) const
        {
            if (_index == 0)
                return n1._point.getX() < n2._point.getX();
            else if (_index == 1)
                return n1._point.getY() < n2._point.getY();

            return false;
        }
        size_t _index;
    };

    TKDNode* makeTree(size_t begin, size_t end, size_t index)
    {
        if (end <= begin)
        {
            return nullptr;
        }

        /* 中间大小的那个位置 */
        size_t n = begin + (end - begin) / 2;
        auto i = m_nodes.begin();
        /* 从小到大排序 */
        std::nth_element(i + begin, i + n, i + end, node_cmp(index));
        index = (index + 1) % 2;
        m_nodes[n]._left  = makeTree(begin, n, index);
        m_nodes[n]._right = makeTree(n + 1, end, index);

        return &m_nodes[n];
    }

    void nearest(TKDNode* root, const CPoint2D& point, size_t index)
    {
        if (root == nullptr)
        {
            return;
        }

        ++m_visited;

        double dist = root->distance(point);

        if (m_best == nullptr || dist < m_best_distance)
        {
            m_best_distance = dist;
            m_best = root;
        }

        if (m_best_distance == 0)
        {
            return;
        }

        double dx = root->get(index) - point.get(index);
        index = (index + 1) % 2;

        nearest(dx > 0 ? root->_left : root->_right, point, index);

        if (dx * dx >= m_best_distance)
        {
            return;
        }

        nearest(dx > 0 ? root->_right : root->_left, point, index);
    }

private:

    TKDNode* m_root = nullptr;
    TKDNode* m_best = nullptr;

    double m_best_distance = 0.0;
    size_t m_visited       = 0;

    std::vector<TKDNode> m_nodes;
};

#endif // CKDTREE_H
